# src/algorithms/stereo_matching.py
import cv2
import numpy as np
from typing import Dict, Optional
from src.data.structures import BinocularImagePair
from src.infrastructure.logger import setupLogger

logger = setupLogger(__name__)


class SGBMStereoMatcher:
    def __init__(self, **params):
        self._params = SGBMStereoMatcher.getDefaultParams()
        self._params.update(params)
        self._matcher = self._create_matcher()

    @staticmethod
    def getDefaultParams() -> Dict[str, int]:
        """获取默认SGBM参数

        Returns:
            Dict[str, int]: 默认SGBM参数字典
        """
        return {
            "min_disparity": 0,
            "num_disparities": 160,
            "block_size": 9,
            "p1_scale": 8,
            "p2_scale": 32,
            "disp12_max_diff": -1,
            "pre_filter_cap": 31,
            "uniqueness_ratio": 10,
            "speckle_window_size": 100,
            "speckle_range": 32,
            "mode_idx": cv2.STEREO_SGBM_MODE_SGBM_3WAY,
        }

    def _create_matcher(self) -> cv2.StereoSGBM:
        """创建 SGBM 匹配器

        Returns:
            cv2.StereoSGBM: SGBM 匹配器实例
        """
        matcher_args = dict(
            minDisparity=self._params["min_disparity"],  # 可能的最小视差值
            numDisparities=self._params[
                "num_disparities"
            ],  # 最大视差减去最小视差, 必须是16的倍数
            blockSize=self._params["block_size"],  # 匹配块大小, 必须为奇数且 >=1
            P1=self._params["p1_scale"]
            * 3
            * self._params["block_size"]
            ** 2,  # 控制视差变化的平滑度, 较小值会产生更多噪声
            P2=self._params["p2_scale"]
            * 3
            * self._params["block_size"] ** 2,  # 控制视差变化的平滑度, 较大值会过度平滑
            disp12MaxDiff=self._params[
                "disp12_max_diff"
            ],  # 左右视差图一致性检查的最大允许差异, 将其设置为非正值可禁用此检查
            preFilterCap=self._params[
                "pre_filter_cap"
            ],  # 预过滤图像像素的截断值. 算法首先计算每个像素的 x 导数, 并将其值裁剪到 [-preFilterCap, preFilterCap] 区间, 结果值被传递给 Birchfield-Tomasi 像素成本函数
            uniquenessRatio=self._params[
                "uniqueness_ratio"
            ],  # 最佳(最小)计算成本函数值"胜过"次佳值的百分比裕度
            speckleWindowSize=self._params[
                "speckle_window_size"
            ],  # 平滑视差区域的最大大小
            speckleRange=self._params["speckle_range"],  # 每个连通分量内的最大视差变化
            mode=[
                cv2.STEREO_SGBM_MODE_SGBM,
                cv2.STEREO_SGBM_MODE_HH,
                cv2.STEREO_SGBM_MODE_SGBM_3WAY,
                cv2.STEREO_SGBM_MODE_HH4,
            ][self._params["mode_idx"]],
        )

        logger.info("SGBM matcher created.")

        return cv2.StereoSGBM.create(**matcher_args)

    def computeDisparity(self, rectified_pair: BinocularImagePair) -> np.ndarray | None:
        """计算视差图

        计算视差时, 超出有效视差范围的值会被设置为 NaN.

        Args:
            rectified_pair (StereoImagePair): 经过校正的立体图像对

        Raises:
            ValueError: 如果输入的图像对无效

        Returns:
            np.ndarray: 视差图
        """
        if rectified_pair.validate() is False:
            logger.error(
                "Invalid StereoImagePair: images are missing or sizes do not match"
            )
            return None

        if rectified_pair.left_image.dtype != np.uint8:
            rectified_pair.left_image = rectified_pair.left_image.astype(np.uint8)
        if rectified_pair.right_image.dtype != np.uint8:
            rectified_pair.right_image = rectified_pair.right_image.astype(np.uint8)

        disparity = (
            self._matcher.compute(
                rectified_pair.left_image, rectified_pair.right_image
            ).astype(np.float32)
            / 16.0
        )

        # 过滤无效视差
        disparity[disparity <= self._params["min_disparity"]] = np.nan
        disparity[
            disparity
            >= (self._params["num_disparities"] + self._params["min_disparity"])
        ] = np.nan

        return disparity

    def updateParams(self, **new_params):
        """更新参数"""
        # 只保留 SGBM 合法键
        valid_keys = {
            "channels",
            "min_disparity",
            "num_disparities",
            "block_size",
            "p1_scale",
            "p2_scale",
            "disp12_max_diff",
            "pre_filter_cap",
            "uniqueness_ratio",
            "speckle_window_size",
            "speckle_range",
            "mode_idx",
        }
        changed = {k: v for k, v in new_params.items() if k in valid_keys}
        logger.debug(f"changed:\n{changed}")

        if not changed:
            return

        # 更新内部字典
        self._params.update(changed)

        # 触发更新
        if "channels" in changed:
            block_size = self._params["block_size"]
            p1_scale = self._params["p1_scale"]
            p2_scale = self._params["p2_scale"]
            self._matcher.setBlockSize(block_size | 1)  # 必须为奇数
            self._matcher.setP1(p1_scale * changed["channels"] * block_size**2)
            self._matcher.setP2(p2_scale * changed["channels"] * block_size**2)
        if "min_disparity" in changed:
            self._matcher.setMinDisparity(changed["min_disparity"])
        if "num_disparities" in changed:
            self._matcher.setNumDisparities(changed["num_disparities"])
        if "block_size" in changed:
            channels = self._params.get("channels", 3)
            p1_scale = self._params["p1_scale"]
            p2_scale = self._params["p2_scale"]
            self._matcher.setBlockSize(changed["block_size"] | 1)  # 必须为奇数
            self._matcher.setP1(p1_scale * channels * changed["block_size"] ** 2)
            self._matcher.setP2(p2_scale * channels * changed["block_size"] ** 2)
        if "p1_scale" in changed:
            channels = self._params.get("channels", 3)
            block_size = self._params["block_size"]
            p1_scale = changed["p1_scale"]
            self._matcher.setBlockSize(block_size | 1)  # 必须为奇数
            self._matcher.setP1(p1_scale * channels * block_size**2)
        if "p2_scale" in changed:
            channels = self._params.get("channels", 3)
            block_size = self._params["block_size"]
            p2_scale = changed["p2_scale"]
            self._matcher.setBlockSize(block_size | 1)  # 必须为奇数
            self._matcher.setP2(p2_scale * channels * block_size**2)
        if "disp12_max_diff" in changed:
            self._matcher.setDisp12MaxDiff(changed["disp12_max_diff"])
        if "pre_filter_cap" in changed:
            self._matcher.setPreFilterCap(changed["pre_filter_cap"])
        if "uniqueness_ratio" in changed:
            self._matcher.setUniquenessRatio(changed["uniqueness_ratio"])
        if "speckle_window_size" in changed:
            self._matcher.setSpeckleWindowSize(changed["speckle_window_size"])
        if "speckle_range" in changed:
            self._matcher.setSpeckleRange(changed["speckle_range"])
        if "mode_idx" in changed:
            self._matcher.setMode(
                [
                    cv2.STEREO_SGBM_MODE_SGBM,
                    cv2.STEREO_SGBM_MODE_HH,
                    cv2.STEREO_SGBM_MODE_SGBM_3WAY,
                    cv2.STEREO_SGBM_MODE_HH4,
                ][self._params["mode_idx"]]
            )

    def postFilter(
        self,
        disparity: np.ndarray,
        left_view: np.ndarray,
        *,
        use_wls: bool = True,
        use_speckle: bool = True,
        use_hierarchy_fill: bool = True,
        wls_lambda: float = 8000.0,
        wls_sigma: float = 1.5,
        speckle_max_size: int = 100,
        speckle_max_diff: float = 1.0,
    ) -> np.ndarray:
        disp = disparity.copy()

        # 1. WLS 边缘保持滤波（小洞闭合 + 去噪）
        if use_wls:
            wls = cv2.ximgproc.createDisparityWLSFilter(self._matcher)
            wls.setLambda(wls_lambda)
            wls.setSigmaColor(wls_sigma)

            disp_wls = np.nan_to_num(disp, nan=0.0)  # WLS 要求 0 代表无效, 而非 NaN

            # 构造占位符, 满足 OpenCV 断言
            h, w = disp_wls.shape[:2]
            disp_right_dummy = np.zeros((h, w), dtype=np.float32)

            disp = wls.filter(disp_wls, left_view, disparity_map_right=disp_right_dummy)

        # 2. Speckle 剔除孤立 blob（可能由 WLS 重新暴露的 0 值）
        if use_speckle:
            cv2.filterSpeckles(
                disp,
                newVal=0,
                maxSpeckleSize=speckle_max_size,
                maxDiff=speckle_max_diff,
            )

        # 3. 积分图多级均值填大洞
        # if use_hierarchy_fill:
        #     mask_invalid = ~np.isfinite(disp) | (disp == 0)
        #     disp = self._hierarchy_mean_fill(disp, mask_invalid)

        return disp

    def _hierarchy_mean_fill(
        self, disp: np.ndarray, invalid: np.ndarray, win_list=None
    ) -> np.ndarray:
        """窗口 21→3 逐级向内均值填充, 不糊边缘"""
        if win_list is None:
            win_list = [21, 15, 9, 5, 3]
        filled = disp.copy()
        for w in win_list:
            intg = cv2.integral(filled.astype(np.float64))
            cnt = cv2.integral((~invalid).astype(np.float64))
            mean = (intg[2:, 2:] - intg[:-2, 2:] - intg[2:, :-2] + intg[:-2, :-2]) / (
                cnt[2:, 2:] - cnt[:-2, 2:] - cnt[2:, :-2] + cnt[:-2, :-2] + 1e-6
            )
            filled[invalid] = mean[invalid]
            invalid = (~np.isfinite(filled)) | (filled == 0)
        return filled
